import json
import random
from typing import Dict, List, Any


class GraphDataGenerator:
    def __init__(self):
        self.colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#F9A602", "#8E44AD",
                       "#2ECC71", "#E74C3C", "#3498DB", "#F39C12", "#1ABC9C"]

        self.node_configs = {
            "自然人": {
                "prefix": "家人_",
                "values": ['赵某', '钱某', '孙某', '李某', '周某', '吴某', '郑某', '王某']
            },
            "手机号码": {
                "prefix": "139",
                "count": 9999999
            },
            "群组": {
                "prefix": "群组",
                "count": 999
            },
            "机顶盒": {
                "prefix": "机顶盒",
                "count": 999
            },
            "夜间常驻基站编码": {
                "prefix": "基站",
                "count": 99
            },
            "小区": {
                "prefix": "",
                "values": ["朝阳区望京新城", "海淀中关村小区", "东城区王府井小区",
                           "西城区金融街小区", "大兴区亦庄小区"]
            },
            "区县街道": {
                "prefix": "",
                "values": ["朝阳区", "海淀区", "东城区", "西城区", "大兴区"]
            }
        }

        self.edge_configs = [
            {"source": "自然人", "target": "手机号码", "label": "拥有"},
            {"source": "手机号码", "target": "手机号码", "label": "通话"},
            {"source": "手机号码", "target": "手机号码", "label": "短信"},
            {"source": "自然人", "target": "夜间常驻基站编码", "label": "夜间常驻"},
            {"source": "自然人", "target": "机顶盒", "label": "拥有"},
            {"source": "自然人", "target": "群组", "label": "归属"},
            {"source": "小区", "target": "区县街道", "label": "位于"}
        ]

        self.vertexSet = {}
        self.edgeSet = {}
        self.groupSet = {}
        self.node_counters = {}
        self.edge_counter = 1
        self.group_counter = 1

    def generate_node_id(self, node_type: str) -> str:
        config = self.node_configs[node_type]
        if "values" in config:
            value = random.choice(config["values"])
            return f"{config['prefix']}{value}"
        else:
            if node_type not in self.node_counters:
                self.node_counters[node_type] = 1
            node_id = f"{config['prefix']}{self.node_counters[node_type]:07d}"
            self.node_counters[node_type] += 1
            return node_id

    def generate_vertex(self, node_type: str, group_id: str) -> Dict[str, Any]:
        node_id = self.generate_node_id(node_type)
        object_key = f"23fsdfw-41f3-11f0{node_id}"

        return {
            "label": node_type,
            "id": node_id,
            "groupId": group_id,
            "properties": {
                "object_key": object_key
            }
        }

    def generate_edge(self, src_id: str, dst_id: str, label: str) -> Dict[str, Any]:
        edge_id = f"edge_{self.edge_counter}"
        self.edge_counter += 1

        return {
            "srcId": src_id,
            "dstId": dst_id,
            "label": label,
            "id": edge_id,
            "dashed": 0,
            "properties": {
                "from_key": f"23fsdfw-41f3-11f0{src_id}",
                "to_key": f"23fsdfw-41f3-11f0{dst_id}",
                "object_key": f"23fsdfw-41f3-11f0{edge_id}"
            }
        }

    def generate_group(self, label: str) -> str:
        group_id = f"g{self.group_counter}"
        color = self.colors[(self.group_counter - 1) % len(self.colors)]

        self.groupSet[group_id] = {
            "label": label,
            "id": group_id,
            "color": color,
            "properties": {}
        }

        self.group_counter += 1
        return group_id

    def generate_data(self, num_nodes_per_type: int = 5, num_edges_per_relation: int = 3) -> Dict[str, Any]:
        # 生成组和节点
        node_groups = {}
        for node_type in self.node_configs.keys():
            group_id = self.generate_group(node_type)
            node_groups[node_type] = group_id

            for _ in range(num_nodes_per_type):
                vertex = self.generate_vertex(node_type, group_id)
                self.vertexSet[vertex["id"]] = vertex

        # 生成边
        for edge_config in self.edge_configs:
            source_type = edge_config["source"]
            target_type = edge_config["target"]
            label = edge_config["label"]

            source_nodes = [node_id for node_id, node in self.vertexSet.items()
                            if node["label"] == source_type]
            target_nodes = [node_id for node_id, node in self.vertexSet.items()
                            if node["label"] == target_type]

            if not source_nodes or not target_nodes:
                continue

            for _ in range(num_edges_per_relation):
                src_id = random.choice(source_nodes)

                # 如果是同一类型的节点，确保不创建自环
                if source_type == target_type:
                    dst_id = random.choice([n for n in target_nodes if n != src_id])
                else:
                    dst_id = random.choice(target_nodes)

                edge = self.generate_edge(src_id, dst_id, label)
                self.edgeSet[edge["id"]] = edge

        return {
            "cypherStatement": "MATCH (u:NATURAL_PERSON) WHERE u.np_louvain_1 = 10215457 match path = (u)-[r*1..3]-(u2:NATURAL_PERSON) where u2.np_louvain_1 = 10215457 return path",
            "resultSet": {
                "header": [],
                "table": []
            },
            "vertexSet": self.vertexSet,
            "edgeSet": self.edgeSet,
            "groupSet": self.groupSet,
            "profile": {
                "totalTime": 9.420103000000001,
                "compileTime": 0.606194,
                "optimizationTime": 1.365633,
                "execgenTime": 0.044288,
                "execTime": 7.348298
            },
            "errors": [
            ],
            "vertexCount": 3,
            "edgeCount": 2,
            "resultCount": 1
        }


# 使用示例
if __name__ == "__main__":
    generator = GraphDataGenerator()
    graph_data = generator.generate_data(num_nodes_per_type=5, num_edges_per_relation=3)

    # 输出JSON数据
    print(json.dumps(graph_data, ensure_ascii=False))